# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: leon

import json
import logging
from typing import List, Union

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.deribit_v1.kr_rest.futures_product import DeribitFuturesProduct


class DeribitSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)

    self._register_url('wss://www.deribit.com/ws/api/v2')
    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)
    self._group_id = group_id

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('deribit_ws_depth_%s%s' % (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('deribit_ws_trade_%s%s' % (group_id, worker_id))[0]
    topic_ticker = topic_map.search_topic_id('deribit_ws_ticker_%s%s' % (group_id, worker_id))[0]
    topic_interest_rate = topic_map.search_topic_id('deribit_ws_interest_rate_%s%s' %
                                                    (group_id, worker_id))[0]

    self._publisher = {
        'book': queue.get_writer(topic_id=topic_depth[1]),
        'trades': queue.get_writer(topic_id=topic_trade[1]),
        'ticker': queue.get_writer(topic_id=topic_ticker[1]),
        'perpetual': queue.get_writer(topic_id=topic_interest_rate[1]),
    }

    self._enable_topic_event_publisher(queue,
                                       'deribit_ws_topic_event', [topic_depth[0], topic_trade[0]])
    self._add_teardown_callback(self._teardown_publisher)

  @gen.coroutine
  def on_open(self):
    yield self._send_subscribe_request()

  @gen.coroutine
  def _send_subscribe_request(self):
    products = [DeribitFuturesProduct.FromStrNativeProduct(symbol) for symbol in self._symbol]
    group_idx = ord(self._group_id) - ord('a')
    id = group_idx
    book_channels = ['book.%s.raw' % product.native_subscription_symbol for product in products]
    subscribe_depth_str = json.dumps({
        "jsonrpc": "2.0",
        "method": "public/subscribe",
        "id": id,
        "params": {
            "channels": book_channels
        }
    })
    yield self._ws.write_message(subscribe_depth_str)
    yield gen.sleep(0.5)

    id += 1
    trade_channels = ['trades.%s.raw' % product.native_subscription_symbol for product in products]
    subscribe_trade_str = json.dumps({
        "jsonrpc": "2.0",
        "method": "public/subscribe",
        "id": id,
        "params": {
            "channels": trade_channels
        }
    })
    yield self._ws.write_message(subscribe_trade_str)
    yield gen.sleep(0.5)

    id += 1
    ticker_channels = [
        'ticker.%s.100ms' % product.native_subscription_symbol for product in products
    ]
    subscribe_ticker_str = json.dumps({
        "jsonrpc": "2.0",
        "method": "public/subscribe",
        "id": id,
        "params": {
            "channels": ticker_channels
        }
    })
    yield self._ws.write_message(subscribe_ticker_str)
    yield gen.sleep(0.5)

    id += 1
    perpetual_products = [product for product in products if product.contract_type == 'perpetual']
    interest_rate_channels = [
        'perpetual.%s.100ms' % product.native_subscription_symbol for product in perpetual_products
    ]
    subscribe_interest_rate_str = json.dumps({
        "jsonrpc": "2.0",
        "method": "public/subscribe",
        "id": id,
        "params": {
            "channels": interest_rate_channels
        }
    })
    yield self._ws.write_message(subscribe_interest_rate_str)

    self._publish_topic_event('START')
    self._logger.info('Subscribe done')

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: bytes):
    if self.stopped:
      return

    if 'id' in message:
      id = message['id']
      group_idx = ord(self._group_id) - ord('a')
      assert group_idx <= id <= group_idx + 3
      self._publish_topic_event(message)
    else:
      try:
        channel = message['params']['channel'].split('.')[0]
        self._publisher[channel].write(timestamp, message)
      except (KeyError, ValueError):
        self._logger.error('Unknown channel: %s' % message)
